{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json, os, pickle\n",
    "from tqdm import tqdm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# obj_categories = ['adult', 'car', 'guitar', 'chair', 'handbag', 'toy', 'baby_seat', 'cat', 'bottle', 'backpack', 'motorcycle', 'ball/sports_ball', 'laptop', 'table', 'surfboard', 'camera', 'sofa', 'screen/monitor', 'bicycle', 'vegetables', 'dog', 'fruits', 'cake', 'cellphone', 'cup', 'bench', 'snowboard', 'skateboard', 'bread', 'bus/truck', 'ski', 'suitcase', 'stool', 'bat', 'elephant', 'fish', 'baby_walker', 'dish', 'watercraft', 'scooter', 'pig', 'refrigerator', 'horse', 'crab', 'bird', 'piano', 'cattle/cow', 'lion', 'chicken', 'camel', 'electric_fan', 'toilet', 'sheep/goat', 'rabbit', 'train', 'penguin', 'hamster/rat', 'snake', 'frisbee', 'aircraft', 'oven', 'racket', 'faucet', 'antelope', 'duck', 'stop_sign', 'sink', 'kangaroo', 'stingray', 'turtle', 'tiger', 'crocodile', 'bear', 'microwave', 'traffic_light', 'panda', 'leopard', 'squirrel']\n",
    "with open('obj_categories.json', 'r') as f:\n",
    "    obj_categories = json.load(f)\n",
    "\n",
    "# obj_to_idx = {}\n",
    "# for i, obj in enumerate(obj_categories):\n",
    "#     obj_to_idx[obj] = i\n",
    "# idx_to_obj = {v:k for k, v in obj_to_idx.items()}\n",
    "# print(len(obj_categories))\n",
    "\n",
    "with open('obj_to_idx.pkl', 'rb') as f:\n",
    "    obj_to_idx = pickle.load(f)\n",
    "with open('idx_to_obj.pkl', 'rb') as f:\n",
    "    idx_to_obj = pickle.load(f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 131/131 [00:08<00:00, 16.31it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "nb_of_videos: 835\n",
      "nb_of_frames: 22976\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "human_categories = ['adult', 'child', 'baby']\n",
    "\n",
    "def convert_vidor_to_detectron2_label(annot_dir, small_dataset=False, val=False):\n",
    "    frame_annots = []\n",
    "    if val:\n",
    "        with open('keyframes_set.pkl', 'rb') as f:\n",
    "            keyframes_set = pickle.load(f)\n",
    "    \n",
    "    nb_of_videos = nb_of_frames = 0\n",
    "    for folder in tqdm(os.listdir(annot_dir)):\n",
    "        for video_json in os.listdir(os.path.join(annot_dir, folder)):\n",
    "            nb_of_videos += 1\n",
    "            with open(os.path.join(annot_dir, folder, video_json), 'r') as f:\n",
    "                annot = json.load(f)\n",
    "\n",
    "            if small_dataset:\n",
    "                frame_per_video_count = 0\n",
    "            for i, frame_label in enumerate(annot['trajectories']):\n",
    "                objs = []\n",
    "                \n",
    "                contains_at_least_one_gt_box = False\n",
    "                # disable contains_at_least_one_gt_box when generating validation frames \n",
    "                # (need to find keyframes instead)\n",
    "                if val: \n",
    "                    contains_at_least_one_gt_box = True\n",
    "                for obj in frame_label:\n",
    "                    if obj['generated'] == 0: # only pick manually labeled frames for training instance detector\n",
    "                        contains_at_least_one_gt_box = True\n",
    "                    label = obj\n",
    "                    label['object_class'] = obj_to_idx['person'] if annot['subject/objects'][obj['tid']]['category'] in human_categories else obj_to_idx[annot['subject/objects'][obj['tid']]['category']]\n",
    "                    objs.append(label)\n",
    "                if not contains_at_least_one_gt_box:\n",
    "                    continue\n",
    "                \n",
    "                if val:\n",
    "                    if (folder + '/' + annot['video_id'], i) not in keyframes_set:\n",
    "                        continue\n",
    "                        \n",
    "                frame_annots.append({ # 'image_id': annot['video_id'] + '/' + str(idx),\n",
    "                    'video_folder': folder,\n",
    "                    'video_id': annot['video_id'],\n",
    "                    'frame_id': str(f'{i+1:06d}'), # 1-based index (ava-style)\n",
    "                    'video_fps': annot['fps'],\n",
    "                    'height': annot['height'],\n",
    "                    'width': annot['width'],\n",
    "                    'middle_frame_timestamp': i,\n",
    "                    'objs': objs,\n",
    "                })\n",
    "                nb_of_frames += 1\n",
    "                \n",
    "                if not val and small_dataset:\n",
    "                    frame_per_video_count += 1\n",
    "                    if frame_per_video_count == 11:\n",
    "                        break\n",
    "\n",
    "    print('nb_of_videos:', nb_of_videos)\n",
    "    print('nb_of_frames:', nb_of_frames)\n",
    "    return frame_annots\n",
    "        \n",
    "train_annot_dir = 'annotation/training'\n",
    "val_annot_dir = 'annotation/validation'\n",
    "\n",
    "# train_frame_annots = convert_vidor_to_detectron2_label(train_annot_dir)\n",
    "# val_frame_annots = convert_vidor_to_detectron2_label(val_annot_dir)\n",
    "\n",
    "# train_frame_annots = convert_vidor_to_detectron2_label(train_annot_dir, small_dataset=True)\n",
    "val_frame_annots = convert_vidor_to_detectron2_label(val_annot_dir, small_dataset=True, val=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[{'video_folder': '0085',\n",
       "  'video_id': '5018581116',\n",
       "  'frame_id': '000011',\n",
       "  'video_fps': 15,\n",
       "  'height': 240,\n",
       "  'width': 320,\n",
       "  'middle_frame_timestamp': 10,\n",
       "  'objs': [{'tid': 0,\n",
       "    'bbox': {'xmin': 87, 'ymin': 2, 'xmax': 206, 'ymax': 98},\n",
       "    'generated': 1,\n",
       "    'tracker': 'mosse',\n",
       "    'object_class': 17},\n",
       "   {'tid': 1,\n",
       "    'bbox': {'xmin': 205, 'ymin': 12, 'xmax': 319, 'ymax': 171},\n",
       "    'generated': 1,\n",
       "    'tracker': 'mosse',\n",
       "    'object_class': 17}]},\n",
       " {'video_folder': '0085',\n",
       "  'video_id': '5018581116',\n",
       "  'frame_id': '000012',\n",
       "  'video_fps': 15,\n",
       "  'height': 240,\n",
       "  'width': 320,\n",
       "  'middle_frame_timestamp': 11,\n",
       "  'objs': [{'tid': 0,\n",
       "    'bbox': {'xmin': 86, 'ymin': 2, 'xmax': 206, 'ymax': 100},\n",
       "    'generated': 1,\n",
       "    'tracker': 'mosse',\n",
       "    'object_class': 17},\n",
       "   {'tid': 1,\n",
       "    'bbox': {'xmin': 206, 'ymin': 13, 'xmax': 319, 'ymax': 172},\n",
       "    'generated': 1,\n",
       "    'tracker': 'mosse',\n",
       "    'object_class': 17}]},\n",
       " {'video_folder': '0085',\n",
       "  'video_id': '5018581116',\n",
       "  'frame_id': '000013',\n",
       "  'video_fps': 15,\n",
       "  'height': 240,\n",
       "  'width': 320,\n",
       "  'middle_frame_timestamp': 12,\n",
       "  'objs': [{'tid': 0,\n",
       "    'bbox': {'xmin': 86, 'ymin': 2, 'xmax': 205, 'ymax': 101},\n",
       "    'generated': 1,\n",
       "    'tracker': 'mosse',\n",
       "    'object_class': 17},\n",
       "   {'tid': 1,\n",
       "    'bbox': {'xmin': 207, 'ymin': 14, 'xmax': 319, 'ymax': 173},\n",
       "    'generated': 1,\n",
       "    'tracker': 'mosse',\n",
       "    'object_class': 17}]},\n",
       " {'video_folder': '0085',\n",
       "  'video_id': '5018581116',\n",
       "  'frame_id': '000018',\n",
       "  'video_fps': 15,\n",
       "  'height': 240,\n",
       "  'width': 320,\n",
       "  'middle_frame_timestamp': 17,\n",
       "  'objs': [{'tid': 0,\n",
       "    'bbox': {'xmin': 83, 'ymin': 0, 'xmax': 202, 'ymax': 107},\n",
       "    'generated': 1,\n",
       "    'tracker': 'mosse',\n",
       "    'object_class': 17},\n",
       "   {'tid': 1,\n",
       "    'bbox': {'xmin': 205, 'ymin': 20, 'xmax': 319, 'ymax': 179},\n",
       "    'generated': 1,\n",
       "    'tracker': 'mosse',\n",
       "    'object_class': 17}]},\n",
       " {'video_folder': '0085',\n",
       "  'video_id': '5018581116',\n",
       "  'frame_id': '000019',\n",
       "  'video_fps': 15,\n",
       "  'height': 240,\n",
       "  'width': 320,\n",
       "  'middle_frame_timestamp': 18,\n",
       "  'objs': [{'tid': 0,\n",
       "    'bbox': {'xmin': 82, 'ymin': 0, 'xmax': 201, 'ymax': 107},\n",
       "    'generated': 1,\n",
       "    'tracker': 'mosse',\n",
       "    'object_class': 17},\n",
       "   {'tid': 1,\n",
       "    'bbox': {'xmin': 204, 'ymin': 21, 'xmax': 319, 'ymax': 180},\n",
       "    'generated': 1,\n",
       "    'tracker': 'mosse',\n",
       "    'object_class': 17}]},\n",
       " {'video_folder': '0085',\n",
       "  'video_id': '3851913128',\n",
       "  'frame_id': '000002',\n",
       "  'video_fps': 29.97002997002997,\n",
       "  'height': 480,\n",
       "  'width': 640,\n",
       "  'middle_frame_timestamp': 1,\n",
       "  'objs': [{'tid': 0,\n",
       "    'bbox': {'xmin': 136, 'ymin': 96, 'xmax': 618, 'ymax': 479},\n",
       "    'generated': 1,\n",
       "    'tracker': 'kcf',\n",
       "    'object_class': 0},\n",
       "   {'tid': 1,\n",
       "    'bbox': {'xmin': 127, 'ymin': 92, 'xmax': 286, 'ymax': 361},\n",
       "    'generated': 1,\n",
       "    'tracker': 'mosse',\n",
       "    'object_class': 0},\n",
       "   {'tid': 2,\n",
       "    'bbox': {'xmin': 0, 'ymin': 165, 'xmax': 138, 'ymax': 361},\n",
       "    'generated': 1,\n",
       "    'tracker': 'kcf',\n",
       "    'object_class': 0}]},\n",
       " {'video_folder': '0085',\n",
       "  'video_id': '3851913128',\n",
       "  'frame_id': '000003',\n",
       "  'video_fps': 29.97002997002997,\n",
       "  'height': 480,\n",
       "  'width': 640,\n",
       "  'middle_frame_timestamp': 2,\n",
       "  'objs': [{'tid': 0,\n",
       "    'bbox': {'xmin': 135, 'ymin': 96, 'xmax': 617, 'ymax': 479},\n",
       "    'generated': 1,\n",
       "    'tracker': 'kcf',\n",
       "    'object_class': 0},\n",
       "   {'tid': 1,\n",
       "    'bbox': {'xmin': 125, 'ymin': 92, 'xmax': 284, 'ymax': 361},\n",
       "    'generated': 1,\n",
       "    'tracker': 'mosse',\n",
       "    'object_class': 0},\n",
       "   {'tid': 2,\n",
       "    'bbox': {'xmin': 0, 'ymin': 166, 'xmax': 137, 'ymax': 362},\n",
       "    'generated': 1,\n",
       "    'tracker': 'kcf',\n",
       "    'object_class': 0}]},\n",
       " {'video_folder': '0085',\n",
       "  'video_id': '3851913128',\n",
       "  'frame_id': '000004',\n",
       "  'video_fps': 29.97002997002997,\n",
       "  'height': 480,\n",
       "  'width': 640,\n",
       "  'middle_frame_timestamp': 3,\n",
       "  'objs': [{'tid': 0,\n",
       "    'bbox': {'xmin': 132, 'ymin': 98, 'xmax': 614, 'ymax': 479},\n",
       "    'generated': 1,\n",
       "    'tracker': 'kcf',\n",
       "    'object_class': 0},\n",
       "   {'tid': 1,\n",
       "    'bbox': {'xmin': 123, 'ymin': 93, 'xmax': 282, 'ymax': 362},\n",
       "    'generated': 1,\n",
       "    'tracker': 'mosse',\n",
       "    'object_class': 0},\n",
       "   {'tid': 2,\n",
       "    'bbox': {'xmin': 0, 'ymin': 167, 'xmax': 136, 'ymax': 363},\n",
       "    'generated': 1,\n",
       "    'tracker': 'kcf',\n",
       "    'object_class': 0}]},\n",
       " {'video_folder': '0085',\n",
       "  'video_id': '3851913128',\n",
       "  'frame_id': '000005',\n",
       "  'video_fps': 29.97002997002997,\n",
       "  'height': 480,\n",
       "  'width': 640,\n",
       "  'middle_frame_timestamp': 4,\n",
       "  'objs': [{'tid': 0,\n",
       "    'bbox': {'xmin': 131, 'ymin': 98, 'xmax': 613, 'ymax': 479},\n",
       "    'generated': 1,\n",
       "    'tracker': 'kcf',\n",
       "    'object_class': 0},\n",
       "   {'tid': 1,\n",
       "    'bbox': {'xmin': 120, 'ymin': 92, 'xmax': 279, 'ymax': 361},\n",
       "    'generated': 1,\n",
       "    'tracker': 'mosse',\n",
       "    'object_class': 0},\n",
       "   {'tid': 2,\n",
       "    'bbox': {'xmin': 0, 'ymin': 168, 'xmax': 135, 'ymax': 364},\n",
       "    'generated': 1,\n",
       "    'tracker': 'kcf',\n",
       "    'object_class': 0}]},\n",
       " {'video_folder': '0085',\n",
       "  'video_id': '3851913128',\n",
       "  'frame_id': '000006',\n",
       "  'video_fps': 29.97002997002997,\n",
       "  'height': 480,\n",
       "  'width': 640,\n",
       "  'middle_frame_timestamp': 5,\n",
       "  'objs': [{'tid': 0,\n",
       "    'bbox': {'xmin': 128, 'ymin': 98, 'xmax': 610, 'ymax': 479},\n",
       "    'generated': 1,\n",
       "    'tracker': 'kcf',\n",
       "    'object_class': 0},\n",
       "   {'tid': 1,\n",
       "    'bbox': {'xmin': 118, 'ymin': 91, 'xmax': 277, 'ymax': 360},\n",
       "    'generated': 1,\n",
       "    'tracker': 'mosse',\n",
       "    'object_class': 0},\n",
       "   {'tid': 2,\n",
       "    'bbox': {'xmin': 0, 'ymin': 167, 'xmax': 132, 'ymax': 363},\n",
       "    'generated': 1,\n",
       "    'tracker': 'kcf',\n",
       "    'object_class': 0}]}]"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "val_frame_annots[:10]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "# with open('train_frame_annots_detectron2.json', 'w') as f:\n",
    "#     json.dump(train_frame_annots, f)\n",
    "with open('val_frame_annots_detectron2.json', 'w') as f:\n",
    "    json.dump(val_frame_annots, f)\n",
    "\n",
    "# with open('train_frame_annots_detectron2_small_dataset_10imgs.json', 'w') as f:\n",
    "#     json.dump(train_frame_annots, f)\n",
    "# with open('val_frame_annots_detectron2_small_dataset_10imgs.json', 'w') as f:\n",
    "#     json.dump(val_frame_annots, f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "22976\n"
     ]
    }
   ],
   "source": [
    "# print(len(train_frame_annots))\n",
    "print(len(val_frame_annots))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "slowfast",
   "language": "python",
   "name": "slowfast"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
